import torch.nn as nn
import torch.nn.functional as F
import torch

from Loss.LossRegister import DefaultLossRegister as LossRegister

@LossRegister.register("GFNet_loss")
class GFNet_loss(nn.Module):
    def __init__(self, alpha = 1, beta = 1):
        super().__init__()
        self.alpha = alpha
        self.beta = beta

    def forward(self, outputs, label):
        x_input = outputs[0]
        x1_input = outputs[1] #GAN输入图片
        x_blur = outputs[2] #虚假图片
        pred = outputs[3] #结果预测
        x_kernel = outputs[4]
        return F.mse_loss(pred,label) * self.alpha + F.mse_loss(x_input, x_blur) * self.beta \
        + torch.norm(x_kernel, 2) + torch.norm(x_blur, 2)


if __name__ == '__main__':
    import torch
    losser = GFNet_loss()
    x = (torch.randn(2,3,224,224),torch.randn(2,3,224,224),torch.randn(2,3,224,224),torch.randn(2,1),torch.randn(2,512,14,14))
    y = torch.randn(2,1)
    x_kernel = torch.randn(2,512,14,14)
    print(torch.norm(x_kernel,1))
    print(losser(x, y))